package com.example.redis.rateLimiter.limiter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.util.Collections;

@Slf4j
public class TokenBucketRateLimiter extends RateLimiter {

    /**
     * redis的lua脚本
     */
    private DefaultRedisScript<Boolean> script;

    /**
     * redisTemplate
     */
    private RedisTemplate<String, Object> redisTemplate;

    public TokenBucketRateLimiter(DefaultRedisScript<Boolean> script, RedisTemplate<String, Object> redisTemplate) {
        this.script = script;
        this.redisTemplate = redisTemplate;
    }

    /**
     * 限流检测(单个接口)
     * @param interfaceName 需要限流的接口名
     * @param maxPermits 最大令牌数
     * @param tokensPerSeconds 每秒生成的令牌数
     * @return 是否通过限流 true: 通过
     */
    @Override
    protected boolean acquire(String interfaceName, long maxPermits, long tokensPerSeconds) {
        // 错误的参数将不起作用
        if (maxPermits <= 0 || tokensPerSeconds <= 0) {
            log.warn("maxPermits and tokensPerSeconds can not be less than zero...");
            return true;
        }

        // 参数结构: KEYS = [限流的key]   ARGV = [最大令牌数, 每秒生成的令牌数, 本次请求的毫秒数]
        Boolean result = this.redisTemplate.execute(this.script, Collections.singletonList(interfaceName), maxPermits, tokensPerSeconds, System.currentTimeMillis());
        return result!=null && result;
    }
}
